#include <my_astar_planner/a_star_planner.h>  // 更新头文件路径
#include <pluginlib/class_list_macros.h>
#include <queue>
#include <unordered_set>

PLUGINLIB_EXPORT_CLASS(my_astar_planner::AStarPlanner, nav_core::BaseGlobalPlanner)

namespace my_astar_planner {

AStarPlanner::AStarPlanner() : costmap_ros_(nullptr), initialized_(false) {}

AStarPlanner::AStarPlanner(std::string name, costmap_2d::Costmap2DROS* costmap_ros) 
    : costmap_ros_(nullptr), initialized_(false) {
    initialize(name, costmap_ros);
}

void AStarPlanner::initialize(std::string name, costmap_2d::Costmap2DROS* costmap_ros) {
    if (!initialized_) {
        costmap_ros_ = costmap_ros;
        costmap_ = costmap_ros_->getCostmap();
        resolution_ = costmap_->getResolution();
        width_ = costmap_->getSizeInCellsX();
        height_ = costmap_->getSizeInCellsY();
        
        ros::NodeHandle private_nh("~/" + name);
        plan_pub_ = private_nh.advertise<nav_msgs::Path>("global_plan", 1);
        
        initialized_ = true;
        ROS_INFO("A* global planner initialized successfully");
    } else {
        ROS_WARN("A* global planner has already been initialized");
    }
}

bool AStarPlanner::makePlan(const geometry_msgs::PoseStamped& start,
                           const geometry_msgs::PoseStamped& goal,
                           std::vector<geometry_msgs::PoseStamped>& plan) {
    if (!initialized_) {
        ROS_ERROR("Planner has not been initialized");
        return false;
    }

    plan.clear();

    // Convert start and goal to map coordinates
    int start_x, start_y, goal_x, goal_y;
    worldToMap(start.pose.position.x, start.pose.position.y, start_x, start_y);
    worldToMap(goal.pose.position.x, goal.pose.position.y, goal_x, goal_y);

    // Check if start and goal are valid
    if (!isValid(start_x, start_y) || !isValid(goal_x, goal_y)) {
        ROS_WARN("Start or goal is outside of the map");
        return false;
    }

    if (!isFree(start_x, start_y) || !isFree(goal_x, goal_y)) {
        ROS_WARN("Start or goal is in occupied space");
        return false;
    }

    // Priority queue for open nodes
    std::priority_queue<Node> open_set;
    
    // Hash sets for open and closed nodes
    std::unordered_set<Node*, NodeHash, NodeEqual> open_set_hash;
    std::unordered_set<Node*, NodeHash, NodeEqual> closed_set;
    
    // Create start node
    double h = heuristic(start_x, start_y, goal_x, goal_y);
    Node* start_node = new Node(start_x, start_y, 0.0, h);
    open_set.push(*start_node);
    open_set_hash.insert(start_node);

    bool found_path = false;
    Node* current_node = nullptr;

    // A* algorithm main loop
    while (!open_set.empty()) {
        // Get node with lowest f value
        current_node = new Node(open_set.top());
        open_set.pop();
        open_set_hash.erase(current_node);

        // Check if we've reached the goal
        if (current_node->x == goal_x && current_node->y == goal_y) {
            found_path = true;
            break;
        }

        closed_set.insert(current_node);

        // Explore neighbors (8-connected grid)
        for (int dx = -1; dx <= 1; ++dx) {
            for (int dy = -1; dy <= 1; ++dy) {
                if (dx == 0 && dy == 0) continue; // Skip current node

                int nx = current_node->x + dx;
                int ny = current_node->y + dy;

                // Check if neighbor is valid and free
                if (!isValid(nx, ny) || !isFree(nx, ny)) {
                    continue;
                }

                // Calculate cost (diagonal moves cost more)
                double move_cost = (dx == 0 || dy == 0) ? 1.0 : 1.414;
                double g = current_node->g + move_cost;
                double h = heuristic(nx, ny, goal_x, goal_y);

                Node* neighbor = new Node(nx, ny, g, h, current_node);

                // Check if neighbor is in closed set
                if (closed_set.find(neighbor) != closed_set.end()) {
                    delete neighbor;
                    continue;
                }

                // Check if neighbor is in open set
                auto it = open_set_hash.find(neighbor);
                if (it != open_set_hash.end()) {
                    // If this path is better, update the node
                    if (g < (*it)->g) {
                        (*it)->g = g;
                        (*it)->f = g + (*it)->h;
                        (*it)->parent = current_node;
                    }
                    delete neighbor;
                } else {
                    // Add to open set
                    open_set.push(*neighbor);
                    open_set_hash.insert(neighbor);
                }
            }
        }
    }

    // Reconstruct path if found
    if (found_path) {
        plan = getPath(current_node);
        publishPlan(plan);
    }

    // Clean up memory
    for (Node* node : open_set_hash) delete node;
    for (Node* node : closed_set) delete node;
    if (!found_path && current_node) delete current_node;

    return found_path;
}

std::vector<geometry_msgs::PoseStamped> AStarPlanner::getPath(Node* node) {
    std::vector<geometry_msgs::PoseStamped> path;
    
    // Backtrack from goal to start
    while (node != nullptr) {
        geometry_msgs::PoseStamped pose;
        mapToWorld(node->x, node->y, pose.pose.position.x, pose.pose.position.y);
        pose.pose.orientation.w = 1.0; // Default orientation
        path.push_back(pose);
        node = node->parent;
    }
    
    // Reverse to get start-to-goal order
    std::reverse(path.begin(), path.end());
    
    return path;
}

double AStarPlanner::heuristic(int x1, int y1, int x2, int y2) {
    // Euclidean distance heuristic
    double dx = x1 - x2;
    double dy = y1 - y2;
    return std::sqrt(dx*dx + dy*dy);
}

bool AStarPlanner::isValid(int x, int y) {
    return x >= 0 && x < width_ && y >= 0 && y < height_;
}

bool AStarPlanner::isFree(int x, int y) {
    unsigned char cost = costmap_->getCost(x, y);
    return cost < costmap_2d::INSCRIBED_INFLATED_OBSTACLE;
}

void AStarPlanner::worldToMap(double wx, double wy, int& mx, int& my) {
    mx = (wx - costmap_->getOriginX()) / resolution_;
    my = (wy - costmap_->getOriginY()) / resolution_;
}

void AStarPlanner::mapToWorld(int mx, int my, double& wx, double& wy) {
    wx = costmap_->getOriginX() + (mx + 0.5) * resolution_;
    wy = costmap_->getOriginY() + (my + 0.5) * resolution_;
}

void AStarPlanner::publishPlan(const std::vector<geometry_msgs::PoseStamped>& path) {
    nav_msgs::Path gui_path;
    gui_path.poses.resize(path.size());
    gui_path.header.frame_id = costmap_ros_->getGlobalFrameID();
    gui_path.header.stamp = ros::Time::now();
    
    for (unsigned int i = 0; i < path.size(); i++) {
        gui_path.poses[i] = path[i];
    }
    
    plan_pub_.publish(gui_path);
}

} // namespace a_star_planner